"""
此文件逻辑是用于训练模型，训练模型的输入可以是普通的csv文件，也可以是ncodecsv文件，训练好的模型会保存到model_path目录下，训练信息会保存到info_path目录下。

"""

import os  
import json  
import matplotlib.pyplot as plt  
import keras  
from keras.layers import Dense, LSTM  
from keras.callbacks import ModelCheckpoint, Callback  
from keras.models import load_model  
from dataset import dataset_features_labels_for_train_test as dataset  
from datetime import datetime
import pandas as pd
import tf2onnx
import tensorflow as tf
from keras.layers import Dropout, BatchNormalization
# 自定义回调：保存最优模型和信息
class SaveBestEpoch(Callback):  
    def __init__(self, save_path, info_path, headers):  
        super().__init__()  
        self.save_path = save_path  
        self.info_path =os.path.join(info_path , f'training_summary.csv')  
        self.headers = headers  # 用于在打印信息时调用  
        self.best_val_loss = float('inf')  
        self.best_epoch = -1
        

    def on_epoch_end(self, epoch, logs=None):  
        current_val_loss = logs.get('val_loss')  
        if current_val_loss < self.best_val_loss:  
            self.best_val_loss = current_val_loss  
            self.best_epoch = epoch + 1  
            # 保存模型  
            self.model.save(self.save_path)  
                   
            # 打印信息  
            print(f"新最佳模型：第 {self.best_epoch} 轮，验证损失：{self.best_val_loss:.6f}")

    def on_train_end(self, logs=None):  
        # 训练结束时，判断是否优于历史记录
        is_better = True
        tag="预测标签"
        if os.path.exists(self.info_path):
            try:
                df = pd.read_csv(self.info_path)
                same_chan = df[df['chan'] == tag]
                if not same_chan.empty:
                    old_best = same_chan['val_loss'].min()
                    if self.best_val_loss >= old_best:
                        print(f" 最优 val_loss={self.best_val_loss:.6f} 未超越历史最优({old_best:.6f})，不保存模型和记录")
                        is_better = False
            except Exception as e:
                print(f"读取历史记录失败，将视为无历史: {e}")

        if is_better:
            # 保存模型
            self.model.save(self.save_path, overwrite=True)
            print(f"训练结束，保存新最佳模型: epoch {self.best_epoch}, val_loss {self.best_val_loss:.6f}")

            # 追加记录
            record = {
                'chan': tag,
                'epoch': self.best_epoch,
                'val_loss': self.best_val_loss,
                'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }
            df_new = pd.DataFrame([record])
            file_exists = os.path.exists(self.info_path)
            df_new.to_csv(self.info_path, mode='a', header=not file_exists, index=False, encoding='utf-8')
            # print(f"{self.chan} 新记录已追加到 CSV")

class LSTMTrainer_V2:  
    def __init__(self, dataset_original,feature_names,label_names, model_path, info_path,train_size=0.8,timestep=50,scaler_path=None):  
        self.feature_names = feature_names
        self.label_names = label_names
               
        self.info_path = info_path  
        self.train_size = train_size  
        self.timestep = timestep
        self.scaler_path = scaler_path  
        self.dataset = dataset.LSTM_Features_Labels_Dataset_Train_Test(dataset_original, feature_names=self.feature_names,label_names=self.label_names, train_size=self.train_size, timestep=self.timestep,scaler_path=self.scaler_path)
  
            
        self.x_train, self.y_train = self.dataset.get_timeclycle_data(is_train=True)  
        self.x_test, self.y_test = self.dataset.get_timeclycle_data(is_train=False)  
        self.model_path = os.path.join(model_path, 'all_model.h5')  
        # 设计模型  
        self.model = self._load_or_create_model()
   

    def _load_or_create_model(self):  
        if os.path.exists(self.model_path):  
            print(f"加载已存在的模型，继续训练...")  
            return load_model(self.model_path)  
        else:  
            print(f"创建新模型，开始训练...")  
            model = keras.Sequential()  
            # model.add(LSTM(80, return_sequences=True, activation='relu'))  
            # model.add(LSTM(100, activation='relu'))  
            # model.add(Dense(10, activation='relu'))  
            # model.add(Dense(1))  
            model.add(LSTM(256, return_sequences=True, activation='tanh'))
            model.add(Dropout(0.2))
            model.add(LSTM(128, return_sequences=False, activation='tanh'))
            model.add(BatchNormalization())
            model.add(Dense(64, activation='relu'))
            model.add(Dense(32, activation='relu'))
            model.add(Dense(len(self.label_names))) 
            model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.001), loss='mse')  
            return model  
    # 训练模型  
    def train_V2(self, continue_epochs=20, batch_size=64): 
         # 初始化回调：保存最优模型和信息  
        save_callback = SaveBestEpoch(self.model_path, self.info_path, headers=self.label_names)  
        checkpoint_callback = ModelCheckpoint(self.model_path, monitor='val_loss', save_best_only=True, mode='min')  

        history = self.model.fit(  
            self.x_train,  
            self.y_train,  
            epochs=continue_epochs,  
            batch_size=batch_size,  
            verbose=2,  
            validation_data=(self.x_test, self.y_test),  
            callbacks=[checkpoint_callback, save_callback]  
        )  

        # 保存最终模型  
        self.model.save(self.model_path)  
        # 转换为 ONNX 并保存
        onnx_path = self.model_path.replace(".h5", ".onnx")
        spec = (tf.TensorSpec(self.x_train.shape, tf.float32, name="input"),)

        # print(f"开始导出 ONNX 模型到 {onnx_path} ...")
        model_proto, _ = tf2onnx.convert.from_keras(
            self.model,
            input_signature=spec,
            opset=13,
            output_path=onnx_path
        )
        print("✅ ONNX 模型保存完成")
       

# 使用示例（移除注释即可运行）  
if __name__ == '__main__':  
    pass